Forward tpu-inference MoE kernel knobs and surface MoE padding - #5087
Merged
Conversation
khatwanimohit
requested review from
A9isha,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
igorts-git,
jiangjy1982,
michelle-yooh,
richjames0,
shralex,
shuningjin,
vipannalla,
xibinliu and
zxhe-sean
as code owners
September 1, 2026 16:20
There was a problem hiding this comment.
Code Review
This pull request updates the vLLM adapter to log a warning when padding moe_intermediate_size due to its memory and FLOP costs. It also forwards environment-backed kernel configuration variables from tpu_inference to fused_moe_func in moe.py to ensure consistent kernel behavior when serving MaxText models through vLLM. The review feedback suggests using getattr with fallback defaults when accessing these environment variables to prevent potential AttributeError crashes with older versions of the tpu_inference package.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
khatwanimohit
force-pushed
the
mohit/vllm-moe-forward-kernel-knobs
branch
from
September 1, 2026 18:23
7d00759 to
82364f1
Compare
khatwanimohit
force-pushed
the
mohit/vllm-moe-forward-kernel-knobs
branch
from
September 1, 2026 20:30
82364f1 to
6c57031
Compare
When a MaxText model is served through vLLM, RoutedMoE.fused_moe_matmul calls tpu-inference's fused_moe_func with only the required arguments. tpu-inference's own serving path passes several environment-backed knobs as well, so the two paths run the same kernel with different configurations, and env vars that operators set on the vLLM command line are silently ignored on the MaxText path. Forward the four environment-backed knobs: ENABLE_RS_KERNEL, USE_GMM_FUSED_RS_KERNEL, ONEHOT_MOE_PERMUTE_THRESHOLD and VLLM_MOE_CHUNK_SIZE. ONEHOT_MOE_PERMUTE_THRESHOLD matters beyond tuning: it selects the one-hot permute path over the SparseCore ragged_gather_reduce kernel, which is what lets expert parallelism run at all on TPU generations whose SparseCore has fewer SIMD lanes than that kernel needs. On v5p, expert_parallelism > 1 previously failed with "num_row_partitions=16 must be <= num_simd_lanes=8". Also make two silent sharding surprises visible in the vLLM adapter. Raise the MoE padding message from absl INFO to WARNING and include its cost. vLLM's logging configuration does not surface absl INFO records, so the padding was invisible: serving Qwen3.5-35B-A3B at moe_mlp_tp_size=4 pads moe_intermediate_size 512 -> 1024 and doubles the MoE weights (measured 33.92 GiB/device instead of 16.95 GiB/device), and at moe_mlp_tp_size=8 it quadruples them, with no indication in the logs. Warn when --enable-expert-parallel is passed but the mesh has no expert shards. The native tpu-inference model paths derive use_ep from parallel_config.enable_expert_parallel, while the mesh takes expert parallelism only from additional_config's sharding_strategy and MaxText derives use_ep from the mesh. The same command line therefore runs expert-parallel MoE natively and tensor-parallel MoE under MaxText, which is the configuration that triggers the padding above. The warning points at expert_parallelism as the fix. Bump the post-training pins to pick these up. USE_GMM_FUSED_RS_KERNEL does not exist in the previously pinned tpu-inference (7ecc401e, Jul 27), and neither does the use_gmm_fused_rs_kernel argument of fused_moe_func, so the change above needs a newer commit to import at all: tpu-inference 7ecc401e -> b67ae5f8 (main) vllm 0ba2aa35 -> d626108b (tpu-inference .buildkite/vllm_lkg.version) Both the old and new tpu-inference pin jax==0.11.0, jaxlib==0.11.0 and libtpu==0.0.44, which already match tpu_post_train_overrides.txt, so the jax version in CI is unchanged. The post-training lock is regenerated with seed-env using the same JAX seed commit recorded for it in docs (52d5cb38). Regeneration needs one fix to run: base requirements ask for an unpinned llguidance, which now resolves to 1.8.0, while vLLM requires >=1.7.0,<1.8.0 in both the old and new commits. Bound it to <1.8.0. The functionally required outcome of the regeneration is tokamax>=0.0.13 (up from 0.0.12, where the GMM_v2 kernel lives) and huggingface-hub>=1.29.0; the lock is installed with --resolution=lowest and the GitHub deps with --no-deps, so these floors are exactly what CI installs.
khatwanimohit
force-pushed
the
mohit/vllm-moe-forward-kernel-knobs
branch
from
September 1, 2026 20:47
6c57031 to
76cd098
Compare
igorts-git
approved these changes
Sep 2, 2026
YixuanWang-99
approved these changes
Sep 2, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Two fixes to the vLLM serving path for MoE models.
RoutedMoE.fused_moe_matmulcalls tpu-inference'sfused_moe_funcwithout theenvironment-backed knobs that tpu-inference passes on its own serving path, so
ENABLE_RS_KERNEL,USE_GMM_FUSED_RS_KERNEL,ONEHOT_MOE_PERMUTE_THRESHOLDandVLLM_MOE_CHUNK_SIZEare silently ignored when serving a MaxText model. Forward them.ONEHOT_MOE_PERMUTE_THRESHOLDmatters beyond tuning: it picks the one-hot permute pathover the SparseCore
ragged_gather_reducekernel, which doesn't fit v5p. Before this,any
expert_parallelism > 1died on the first request withAssertionError: num_row_partitions=16 must be <= num_simd_lanes=8.The adapter pads
moe_intermediate_sizeto satisfy GMM_v2's tiling constraint, whichinflates every expert weight (
moe_mlp_tp_sizealways equals--tensor-parallel-size,so TP=4 pads 512->1024 and TP=8 pads 512->2048). It was logged at absl INFO, which vLLM's
logging doesn't surface, so serving Qwen3.5-35B-A3B quietly used 2x the MoE weights with
nothing in the logs. Log it at WARNING with the cost.
No behaviour change with no env vars set, and the padding logic itself is untouched.
Tests
Qwen3.5-35B-A3B on v5p-8, bf16 (v5p has no fp8 MXU). Server:
with
Client, same for every run:
240/240 turns in every run, identical output tokens (37,194), so wall clock is comparable.
Native vLLM (
MODEL_IMPL_TYPE=vllm, same flags) included for reference.+37% tok/s/chip, and the gap to native goes from 1.58x to 1.15x. Decode is now ahead of
native; what's left is turn-1 prefill, since
tensor_parallelism=1leaves the prefill GEMMsunsharded.
TP=2 / EP=2is also padding-free (18.12 GiB/device) and is the obvious follow-up— it was equally unrunnable before this PR.
The padding warning on the before config, where previously nothing was logged:
Note
--enable-expert-parallelon the CLI doesn't affect the JAX mesh — tpu-inference readsexpert parallelism only from
additional_config, which is why it's set there above.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.